#!/usr/bin/env python
# -*- encoding: utf-8 -*-
'''
@File    :   loss.py
@Time    :   2022-06-21 08:17:22
@Author  :   GuoLiuFang
@Version :   0.1
@Contact :   guoliufangking@gmail.com
@License :   (C)Copyright 2018-2022, RandomMatrix
@Desc    :   None
'''
# import other libs
import torch
import torch.nn as nn
from utils import iou
class YoloLoss(nn.Module):
    def __init__(self, S=7, B=2, C=20):
        super(YoloLoss, self).__init__()
        self.S = S
        self.B = B
        self.C = C
        self.lambda_coord = 5
        self.lambda_noobj = 0.5
        self.mse = nn.MSELoss(reduction='sum')

    def forward(self, predictions, targets):
        # predictions = (32, 7 * 7 * 30)
        # targets = (32, 7 * 7 * (20 + (1 + 4))); x y w h
        predictions = predictions.reshape(-1, self.S, self.S, self.C + 5 * self.B)
        # predictions = 32, 7, 7, 30
        # targets = 32, 7, 7, 25
        # gt_mask = targets[..., 20]
        # gt_mask = 32, 7, 7
        gt_mask = targets[..., 20:21]
        # gt_mask = 32, 7, 7, 1
        # input 32, 7, 7, 4
        # output 32, 7, 7, 1
        iou1 = iou(predictions[..., 21:25], targets[..., 21:25])
        iou2 = iou(predictions[..., 26:30], targets[..., 21:25])
        ious = torch.cat([iou1.unsqueeze(0), iou2.unsqueeze(0)], dim=0)
        # iou1.unsqueeze(0) = 32, 7, 7, 1====> 1, 32, 7, 7, 1
        # iou2.unsqueeze(0) = 32, 7, 7, 1====> 1, 32, 7, 7, 1
        # ious ====> 1, 32, 7, 7, 1
        max_ious, iou_mask = torch.max(ious, dim=0)
        # max_ious, keep value, 32, 7, 7, 1
        # iou_mask, keep index, 0, 1, shape 32, 7, 7, 1
        coord_predictions = iou_mask * predictions[..., 26:30] + (1 - iou_mask) * predictions[..., 21:25]
        # iou_mask 32, 7, 7, 1 
        # predicitons[..., 26:30], 32, 7, 7, 4
        # ????===, 32, 7, 7, 4
        coord_predictions = gt_mask * coord_predictions
        # coord_predictions , 32, 7, 7, 4 (x, y, w, h)
        # targets 也遵守同样的操作。
        coord_predictions[..., 2:4] = torch.sqrt(coord_predictions[..., 2:4])
        coord_targets = gt_mask * targets[..., 21:25]
        coord_targets[..., 2:4] = torch.sqrt(coord_targets[..., 2:4])
        # 32, 7, 7, 4 (x, y, w', h')
 
        coord_loss = self.mse(torch.flatten(coord_predictions, end_dim=-2),torch.flatten(coord_targets, end_dim=-2))
        # coord_loss = 32*7*7, 4======> 32,*7*7, 1
        # 1、计算 x y w h的均方差。

        # 2、计算P的问题
        # 2.1、 obj
        # [..., 25] vs [..., 25:26]之间的区别
        # 32, 7, 7   vs    32, 7, 7, 1
        prob_predictions = gt_mask * (iou_mask * predictions[..., 25:26] + (1 - iou_mask) * predictions[..., 20:21])
        prob_targets = gt_mask * targets[..., 20:21]
        obj_prob_loss = self.mse(torch.flatten(prob_predictions, end_dim=-2), torch.flatten(prob_targets, end_dim=-2))
        # 2.2、noobj
        # ground truth noobj 就是0
        # 
        noobj_prob_targets = (1 - gt_mask) * targets[..., 20:21]

        noobj_prob_prediciton_box1 = (1 - gt_mask) * (predictions[..., 20:21])
        noobj_prob_loss = self.mse(torch.flatten(noobj_prob_prediciton_box1, end_dim=-2), torch.flatten(noobj_prob_targets, end_dim=-2))
        noobj_prob_prediciton_box2 = (1 - gt_mask) * (predictions[..., 25:26])
        noobj_prob_loss += self.mse(torch.flatten(noobj_prob_prediciton_box2, end_dim=-2), torch.flatten(noobj_prob_targets, end_dim=-2))
        
        # 3、计算class的问题
        class_targets = gt_mask * targets[..., :20]
        # class_targets = 32, 7, 7, 20
        class_predictions = gt_mask * predictions[..., :20]
        # class_pre = 32, 7, 7, 20
        class_loss = self.mse(torch.flatten(class_targets, end_dim=-2), torch.flatten(class_predictions, end_dim=-2))

        total_loss = self.lambda_coord * coord_loss + obj_prob_loss + self.lambda_noobj * noobj_prob_loss + class_loss
        return total_loss
        # 32 * 7 * 7, 1